'''
Author: daniel
Date: 2024-03-01 15:41:38
LastEditTime: 2024-03-02 13:13:54
LastEditors: daniel
Description: 
FilePath: /mtpsl/jupyters/utils.py
have a nice day
'''
import numpy as np
import torchvision
import cv2

import torch 
import os 

from os.path import join, exists, split





def transform_semantic_to_colorful_semantics(semantic_map):
    classes = {
        0: (0, 0, 0),       # 无标签
        1: (128, 0, 0),     # 墙壁
        2: (0, 128, 0),     # 地板
        3: (128, 128, 0),   # 柜子
        4: (0, 0, 128),     # 床
        5: (128, 0, 128),   # 椅子
        6: (0, 128, 128),   # 沙发
        7: (128, 128, 128), # 表面
        8: (64, 0, 0),      # 窗户
        9: (192, 0, 0),     # 门
        10: (64, 128, 0),   # 架子
        11: (192, 128, 0),  # 屏幕
        12: (64, 0, 128)    # 窗帘
    }

    semantic_image = semantic_map.numpy()
    colored_image = np.zeros((semantic_image.shape[0], semantic_image.shape[1], 3), dtype=np.uint8)
    for class_id, color in classes.items():
        colored_image[semantic_image == class_id] = color
    
    # plt.imshow(colored_image)
    # plt.axis('off')
    # plt.show()
    return colored_image




def tensor2numpy_img(tensor_img):
    return  (tensor_img.squeeze().permute([1,2,0]).cpu().numpy()* 255 ).astype(np.uint8) 
    

def merge_img(img1, img2, alpha =0.6):
    return cv2.addWeighted(img1,alpha, img2 , 1 - alpha, 0)


def show_img(img):
    return torchvision.transforms.functional.to_pil_image(img)
    
def draw_semantics(image,intrinsics, alpha = 0.6):
    merged = merge_img(tensor2numpy_img(image),transform_semantic_to_colorful_semantics(intrinsics),alpha = alpha)
    return torchvision.transforms.functional.to_pil_image(merged)


def draw_painter_semantics(image,intrinsics, alpha = 0.6):
    merged = merge_img(tensor2numpy_img(image),intrinsics,alpha = alpha)
    return torchvision.transforms.functional.to_pil_image(merged)


def normalize(im):
    return (im - im.min())/ (im.max() - im.min())

def draw_depth(image, intrinsics, alpha = 0.6):
    intrinsics = intrinsics.repeat([3,1,1]).permute(1,2,0) 
    merged = merge_img(tensor2numpy_img(image), (normalize(intrinsics) * 255).numpy().astype(np.uint8),alpha = alpha)
    return torchvision.transforms.functional.to_pil_image(merged)

def draw_normal(image, intrinsics, alpha = 0.6):
    intrinsics = intrinsics.permute(1,2,0) 
    merged = merge_img(tensor2numpy_img(image), (normalize(intrinsics) * 255).numpy().astype(np.uint8),alpha = alpha)
    return torchvision.transforms.functional.to_pil_image(merged)


def make_dir(path):
    if not exists(path):
        os.makedirs(path)
